import numpy as np


colors = [
    [10,20,250],
    [20,40,244],
    [30,60,238],
    [40,80,232],
    [50,60,226],
    [59,126,219],
    [78,138,221],
    [97,150,224],
    [116,163,226],
    [135,175,229],
    [155,188,232],
    [154,196,220],
    [152,214,196],
    [153,205,208],
    [152,214,196],
    [151,232,173],
    [215,222,126],
    [234,219, 112],
    [244, 217, 99],
    [250,204,79],
    [247, 180, 45],
    [241,155,0],
    [241, 147, 3],
    [239, 117, 17],
    [238,102,24],
    [238, 88, 31],
    [231,75,26],
    [224, 63, 22],
    [217,51,18],
    [208, 36, 14],
    [194,0,3],
    [181, 1, 9],
    [169,2,16],
    [138, 5, 25],
    [111,0,21],
    [80, 0, 15],
]


def ImGray(im, category):
    scale = 6
    if category == '2t':
        im += 50
        scale = 2

    im *= scale
    im[im > 255] = 255
    return im.astype('uint8')

# 0--0.1mm， light gray
# 0.1--10mm, light green
# 10--25mm,  green
# 25--50mm,  light blue
# 50--100mm, blue
# 100--inf,  purple
def ImColor_tp(im):
    imc = np.zeros((im.shape[0], im.shape[1], 3))
    imc[:, :] = [215, 215, 215]
    imc[im > 0.1] = [152, 251, 152]
    imc[im > 10] = [64, 189, 64]
    imc[im > 25] = [92, 172, 238]
    imc[im > 50] = [0, 0, 255]
    imc[im > 100] = [238, 0, 238]

    return imc.astype('uint8')

def ImColor_tp_thresold(im,th):
    imc = np.zeros((im.shape[0], im.shape[1], 3))
    imc[:, :] = [215, 215, 215]
    if th <= 0.1:
        imc[im > 0.1] = [152, 251, 152]
    if th <= 10:
        imc[im > 10] = [64, 189, 64]
    if th <= 25:
        imc[im > 25] = [92, 172, 238]
    if th <= 50:
        imc[im > 50] = [0, 0, 255]
    if th <= 100:
        imc[im > 100] = [238, 0, 238]

    return imc.astype('uint8')

def ImColor_tp_diff(x,y):
    imc = np.zeros((y.shape[0], y.shape[1], 3))
    imc[:, :] = [150, 150, 150]
    imc[np.where((y <= 0.1) & (x > 0.1))] = [215,215,215]
    imc[np.where((y > 0.1) & (y <= 10) & ((x <= 0.1) | (x > 10)))] = [152, 251, 152]
    imc[np.where((y > 10) & (y <= 25) & ((x <= 10) | (x > 25)))] = [34, 139, 34]
    imc[np.where((y > 25) & (y <= 50) & ((x <= 25) | (x > 50)))] = [92, 172, 238]
    imc[np.where((y > 50) & (y <= 100) & ((x <= 50) | (x > 100)))] = [0, 0, 255]
    imc[np.where((y > 100) & (x <= 100))] = [238, 0, 238]

# def ImColor_tp_diff2(x,y,r):
#     imc = np.zeros((y.shape[0], y.shape[1], 3))
#     imc[:, :] = [215,215,215]
#     imc[np.where((y <= 0.1) & (x <= 0.1) & (r > 0.1))] = [238, 0, 238]
#     imc[np.where((y > 0.1) & (x > 0.1) & (r <= 0.1))] = [238, 0, 238]
#     imc[np.where((y <= 0.1) & (r <= 0.1) & (x > 0.1))] = [34, 139, 34]
#     imc[np.where((y > 0.1) & (r > 0.1) & (x <= 0.1))] = [34, 139, 34]
#     x1 = np.sum((y <= 0.1) & (x <= 0.1) & (r > 0.1))
#     x2 = np.sum((y > 0.1) & (x > 0.1) & (r <= 0.1))
#     r1 = np.sum((y <= 0.1) & (r <= 0.1) & (x > 0.1))
#     r2 = np.sum((y > 0.1) & (r > 0.1) & (x <= 0.1))
#
#     return imc.astype('uint8'), x1+x2, r1+r2

def ImColor_tp_diff2(th,x,y,r):
    imc = np.zeros((y.shape[0], y.shape[1], 3))
    imc[:, :] = [215,215,215]
    imc[np.where((y > th) & (x > th) & (r <= th))] = [180, 0, 180]
    r_fn = np.sum((y > th) & (x > th) & (r <= th))
    imc[np.where((y <= th) & (x <= th) & (r > th))] = [238, 0, 238]
    r_fp = np.sum((y <= th) & (x <= th) & (r > th))
    xb = r_fn + r_fp

    imc[np.where((y > th) & (r > th) & (x <= th))] = [24, 109, 24]
    x_fn = np.sum((y > th) & (r > th) & (x <= th))
    imc[np.where((y <= th) & (r <= th) & (x > th))] = [34, 139, 34]
    x_fp = np.sum((y <= th) & (r <= th) & (x > th))
    rb = x_fn + x_fp

    a1 = np.sum((y > th))

    return imc.astype('uint8'), xb, rb, a1

def ImColor_tm(im):
    imc = np.zeros((im.shape[0], im.shape[1], 3))
    imc[:,:] = [0,0,255]
    for i, l in enumerate(range(-30,42,2)):
        imc[im > l] = colors[i]

    return imc.astype('uint8')

